{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! /usr/bin/env python\n",
    "# coding=utf-8\n",
    "\n",
    "import cv2\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import core.utils as utils\n",
    "import matplotlib.pylab as plt\n",
    "from PIL import Image\n",
    "from core.config import cfg\n",
    "from core.yolov3 import YOLOv3, decode\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL']='2'\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "tf.config.experimental.set_memory_growth(gpus[0], True)\n",
    "\n",
    "INPUT_SIZE   = 608\n",
    "NUM_CLASS    = len(utils.read_class_names(cfg.YOLO.CLASSES))\n",
    "CLASSES      = utils.read_class_names(cfg.YOLO.CLASSES)\n",
    "\n",
    "# Build Model\n",
    "input_layer  = tf.keras.layers.Input([INPUT_SIZE, INPUT_SIZE, 3])\n",
    "feature_maps = YOLOv3(input_layer)\n",
    "\n",
    "bbox_tensors = []\n",
    "for i, fm in enumerate(feature_maps):\n",
    "    bbox_tensor = decode(fm, i)\n",
    "    bbox_tensors.append(bbox_tensor)\n",
    "\n",
    "model = tf.keras.Model(input_layer, bbox_tensors)\n",
    "model.load_weights(\"./yolov3\")\n",
    "\n",
    "image = cv2.imread(\"C:/Users/ZZK/Documents/auto_drive/train_dataset/train_image/000001.jpg\")\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "image_size = image.shape[:2]\n",
    "image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])\n",
    "image_data = image_data[np.newaxis, ...].astype(np.float32)\n",
    "\n",
    "################################显示图片\n",
    "\n",
    "def show_single_image(img_arr):\n",
    "    plt.imshow(img_arr,cmap='binary')\n",
    "    plt.axis('on')\n",
    "    plt.show()\n",
    "\n",
    "pred_bbox = model.predict(image_data)\n",
    "pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]\n",
    "pred_bbox = tf.concat(pred_bbox, axis=0)\n",
    "bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)\n",
    "bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')\n",
    "image = utils.draw_bbox(image, bboxes)\n",
    "image = Image.fromarray(image)\n",
    "image.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "writefile = './docs/result.txt'\n",
    "file_to_write = open(writefile, 'w')\n",
    "\n",
    "for i in range(7000,10000):\n",
    "    image = cv2.imread(\"C:/Users/ZZK/Documents/auto_drive/test/00\"+str(i)+\".jpg\")\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    image_size = image.shape[:2]\n",
    "    image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])\n",
    "    image_data = image_data[np.newaxis, ...].astype(np.float32)\n",
    "\n",
    "    pred_bbox = model.predict(image_data)\n",
    "    pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]\n",
    "    pred_bbox = tf.concat(pred_bbox, axis=0)\n",
    "    bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)\n",
    "    bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')\n",
    "    print('picture:',i,'/9999')\n",
    "    file_to_write.write(\"00\" + str(i) + \".jpg\" + \" \" + \"00\" + str(i) + \".png\" + \" \")\n",
    "    for j in range(len(bboxes)):\n",
    "        tem = str(int(bboxes[j][0]))+\",\"+str(int(bboxes[j][1]))+\",\"+str(int(bboxes[j][2]))+\",\"+str(int(bboxes[j][3]))+\",\"+str(int(bboxes[j][5]))+\",\"+\"%.2f\"%bboxes[j][4] \n",
    "        file_to_write.write(tem + \" \")\n",
    "    file_to_write.write('\\n')     \n",
    "file_to_write.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###############################显示效果\n",
    "\n",
    "for i in range(7011,7020):\n",
    "    image = cv2.imread(\"C:/Users/ZZK/Documents/auto_drive/test/00\"+str(i)+\".jpg\")\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    image_size = image.shape[:2]\n",
    "    image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])\n",
    "    image_data = image_data[np.newaxis, ...].astype(np.float32)\n",
    "\n",
    "    pred_bbox = model.predict(image_data)\n",
    "    pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]\n",
    "    pred_bbox = tf.concat(pred_bbox, axis=0)\n",
    "    bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)\n",
    "    bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')\n",
    "#     for i in range(len(bboxes)):\n",
    "#         print(bboxes[i][4])\n",
    "    image = utils.draw_bbox(image, bboxes)\n",
    "    image = Image.fromarray(image)\n",
    "    image.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "###############################显示训练集/结果\n",
    "\n",
    "num = 0\n",
    "\n",
    "# with open(\"C:/Users/ZZK/Documents/auto_drive/train_dataset/train.txt\", 'r') as f:\n",
    "with open(\"C:/Users/ZZK/auto_drive/docs/0.04.txt\", 'r') as f:\n",
    "    txt = f.readlines()\n",
    "    for lines in txt:\n",
    "        num += 1\n",
    "        if num <= 24:\n",
    "            continue\n",
    "        if num == 30:\n",
    "            break\n",
    "        line = lines.split( )\n",
    "        line_len = len(line)\n",
    "#         image = cv2.imread(\"C:/Users/ZZK/Documents/auto_drive/train_dataset/train_image/\"+line[0])\n",
    "        image = cv2.imread(\"C:/Users/ZZK/Documents/auto_drive/test/\"+line[0])\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        for i in range(2,line_len):  \n",
    "            coor = line[i].split(',')\n",
    "            cv2.rectangle(image,(int(coor[0]),int(coor[1])),(int(coor[2]),int(coor[3])), (255,0,0),1) \n",
    "            cv2.putText(image, coor[4], (int(coor[0]), int(coor[1])-2), cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                        0.5, (0, 0, 0), 1, lineType=cv2.LINE_AA)\n",
    "            cv2.putText(image, coor[5], (int(coor[2]), int(coor[3])-2), cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                        0.5, (0, 0, 0), 1, lineType=cv2.LINE_AA)\n",
    "        cv2.putText(image, line[0], (20, 20), cv2.FONT_HERSHEY_SIMPLEX,0.5, (0, 0, 0), 1, lineType=cv2.LINE_AA)\n",
    "        image = Image.fromarray(image)\n",
    "        image.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###############################单张显示\n",
    "\n",
    "num = 0\n",
    "with open(\"C:/Users/ZZK/Documents/auto_drive/train_dataset/train.txt\", 'r') as f:\n",
    "    txt = f.readlines()\n",
    "    for lines in txt:\n",
    "        num += 1\n",
    "        if num <= 509:\n",
    "            continue\n",
    "        line = lines.split( )\n",
    "        line_len = len(line)\n",
    "        if line_len == 5:\n",
    "            break\n",
    "        image = cv2.imread(\"C:/Users/ZZK/Documents/auto_drive/train_dataset/train_image/\"+line[0])\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        for i in range(2,line_len):  \n",
    "            coor = line[i].split(',')\n",
    "            cv2.rectangle(image,(int(coor[0]),int(coor[1])),(int(coor[2]),int(coor[3])), (255,0,0),1)\n",
    "        show_single_image(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "writefile = './docs/result.txt'\n",
    "file_to_write = open(writefile, 'w')\n",
    "with open(\"C:/Users/ZZK/auto_drive/docs/0.04.txt\", 'r') as f:\n",
    "    txt = f.readlines()\n",
    "    for lines in txt:\n",
    "        line = lines.split( )\n",
    "        line_len = len(line)\n",
    "        print(line[0])\n",
    "        file_to_write.write(line[0]+' '+line[1]+' ')\n",
    "        for i in range(2,line_len):  \n",
    "            coor = line[i].split(',')\n",
    "            if float(coor[5]) >= 0.4:\n",
    "                file_to_write.write(line[i]+' ')\n",
    "        file_to_write.write('\\n')\n",
    "    \n",
    "file_to_write.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
